import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.nn.functional as F
from torch import optim
import numpy as np
import sys
sys.path.append("/home/project/Qtrainer_slim/test/yolov4") 
from models.yolov4_csp import YoloV4CspNet


def trainnet():
    # 定义数据
    inputdata = np.ones((5,3,608,608))
    input_tr = torch.from_numpy(inputdata)
    input_tr = torch.tensor(input_tr , dtype=torch.float32)
    print("input_tr",input_tr)
    targetdata = np.ones((5,3))*100
    # targetdata = np.array([targetdata,targetdata,targetdata])
    target_tr = torch.from_numpy(targetdata)
    target_tr = torch.tensor(target_tr , dtype=torch.float32)
    # 构建网络
    mynet = YoloV4CspNet()
    # 定义优化器
    # optimizer = optim.SGD(mynet.parameters(),lr=0.01,momentum=0.9)
    optimizer = optim.Adam(mynet.parameters(),lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=False)
    # 定义loss
    compute_loss = nn.MSELoss()

    lossList = []
    # 训练，参数优化
    for i in range(10):
        # 梯度清零
        optimizer.zero_grad()
        # 数据正向传播
        outputdata = mynet(input_tr)
        # print("outputdata",outputdata)
        # print("outputdata",outputdata.shape)
        # 计算loss
        # outnp = np.array([outputdata,outputdata,outputdata])
        # outTensor = torch.from_numpy(outnp)
        # outTensor = torch.tensor(outnp,dtype=torch.float32)
        # outTensor = torch.from_numpy([outputdata,outputdata,outputdata])
        # print("outTensor",outTensor)
        # targetTensor =torch.tensor([target_tr,target_tr,target_tr])
        # outputdata = np.array(outputdata)
        # outputdata  = torch.from_numpy(outputdata)
        # print("outputdata ",outputdata )
        # outputdata = torch.tensor(outputdata)
        # target_tr = Variable(target_tr ,requires_grad=True)
        print("target_tr",target_tr)
        print("outputdata",outputdata)
        # loss = 
     
        loss0 = compute_loss(target_tr,outputdata[0])
        loss1 = compute_loss(target_tr,outputdata[1])
        loss2 = compute_loss(target_tr,outputdata[2])

        loss = abs(loss0)+ abs(loss1)+abs(loss2)
        # loss = compute_loss(target_tr,outputdata)
        # loss = 
        print("loss",loss)
        lossList.append(loss)
        # 误差反向传播
        loss.backward()
        # 更新参数
        optimizer.step()
    # torch.save(mynet, 'yolov4_test_all.pth') 
    # torch.save(mynet.state_dict(), 'yolov4_test_dict.pth') 
    # print("lossList",lossList)
   
trainnet()